package com.sheng.leetcode.year2023.month09.day21;

import lombok.var;
import org.junit.Test;

import java.util.*;

/**
 * @author by ls
 * @date 2023/9/21
 * <p>
 * 2603. 收集树中金币<p>
 * <p>
 * 给你一个 n 个节点的无向无根树，节点编号从 0 到 n - 1 。给你整数 n 和一个长度为 n - 1 的二维整数数组 edges ，<p>
 * 其中 edges[i] = [ai, bi] 表示树中节点 ai 和 bi 之间有一条边。再给你一个长度为 n 的数组 coins ，<p>
 * 其中 coins[i] 可能为 0 也可能为 1 ，1 表示节点 i 处有一个金币。<p>
 * 一开始，你需要选择树中任意一个节点出发。你可以执行下述操作任意次：<p>
 * 收集距离当前节点距离为 2 以内的所有金币，或者<p>
 * 移动到树中一个相邻节点。<p>
 * 你需要收集树中所有的金币，并且回到出发节点，请你返回最少经过的边数。<p>
 * 如果你多次经过一条边，每一次经过都会给答案加一。<p>
 * <p>
 * 示例 1：<p>
 * 输入：coins = [1,0,0,0,0,1], edges = [[0,1],[1,2],[2,3],[3,4],[4,5]]<p>
 * 输出：2<p>
 * 解释：从节点 2 出发，收集节点 0 处的金币，移动到节点 3 ，收集节点 5 处的金币，然后移动回节点 2 。<p>
 * <p>
 * 示例 2：<p>
 * 输入：coins = [0,0,0,1,1,0,0,1], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[5,6],[5,7]]<p>
 * 输出：2<p>
 * 解释：从节点 0 出发，收集节点 4 和 3 处的金币，移动到节点 2 处，收集节点 7 处的金币，移动回节点 0 。<p>
 * <p>
 * 提示：<p>
 * n == coins.length<p>
 * 1 <= n <= 3 * 10^4<p>
 * 0 <= coins[i] <= 1<p>
 * edges.length == n - 1<p>
 * edges[i].length == 2<p>
 * 0 <= ai, bi < n<p>
 * ai != bi<p>
 * edges 表示一棵合法的树。<p>
 */
public class LeetCode2603 {

    @Test
    public void test01() {
        int[] coins = {1, 0, 0, 0, 0, 1};
        int[][] edges = {{0, 1}, {1, 2}, {2, 3}, {3, 4}, {4, 5}};
//        int[] coins = {0, 0, 0, 1, 1, 0, 0, 1};
//        int[][] edges = {{0, 1}, {0, 2}, {1, 3}, {1, 4}, {2, 5}, {5, 6}, {5, 7}};
        System.out.println(new Solution().collectTheCoins(coins, edges));
    }
}

class Solution {
    public int collectTheCoins(int[] coins, int[][] edges) {
        /*
          我们先将 edges 中的边转换成邻接表 g，其中 g[i] 表示节点 i 的所有邻接节点，用集合表示。
          接下来我们遍历所有节点，找到 coins[i]=0
          且 g[i] 中只有一个节点的节点（也即是金币为 0 的叶子节点），将其加入队列 q 中。
          然后我们不断地从队列中取出节点，将其从邻接表中删除，
          然后判断其邻接节点是否满足 coins[j]=0
          且 g[j] 中只有一个节点的条件，如果满足则将其加入队列 q 中。循环，直至队列为空。
          经过上述操作后，我们得到了一棵新的树，且树的叶子节点都是金币为 1 的节点。
          然后，我们再删除剩下的两层叶子节点，最终得到的是一棵所有节点都需要被访问的节点，我们只需要统计其边数，乘上 2，即为答案。

          作者：ylb
          链接：https://leetcode.cn/problems/collect-coins-in-a-tree/
          来源：力扣（LeetCode）
          著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
         */
        int n = coins.length;
        Set<Integer>[] g = new Set[n];
        Arrays.setAll(g, k -> new HashSet<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        Deque<Integer> q = new ArrayDeque<>();
        for (int i = 0; i < n; ++i) {
            if (coins[i] == 0 && g[i].size() == 1) {
                q.offer(i);
            }
        }
        while (!q.isEmpty()) {
            int i = q.poll();
            for (int j : g[i]) {
                g[j].remove(i);
                if (coins[j] == 0 && g[j].size() == 1) {
                    q.offer(j);
                }
            }
            g[i].clear();
        }
        for (int k = 0; k < 2; ++k) {
            for (int i = 0; i < n; ++i) {
                if (g[i].size() == 1) {
                    q.offer(i);
                }
            }
            for (int i : q) {
                for (int j : g[i]) {
                    g[j].remove(i);
                }
                g[i].clear();
            }
        }
        int ans = 0;
        for (var e : edges) {
            int a = e[0], b = e[1];
            if (!g[a].isEmpty() && !g[b].isEmpty()) {
                ans += 2;
            }
        }
        return ans;
    }
}
